from collections.abc import Iterable, Mapping
from functools import cached_property
from typing import Any, TypeVar, TypedDict

import numpy as np

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, DatasetWriter, pause, resume
from zkl_aiutils_datasets.utils.weight import parse_item_weight_pairs

AnySample = TypeVar('AnySample')


class RandomGatheringIteratorState(TypedDict):
    iterators_index: tuple[int, ...]
    iterators_state: tuple[Any, ...]
    random_state: Any


class RandomGatheringIterator(DatasetIterator[AnySample]):
    def __init__(self,
        datasets: 'RandomGatheringDataset[AnySample] | DatasetWeightPairs',
        state: RandomGatheringIteratorState | None = None,
        random_seed: int = 42,
    ):
        if isinstance(datasets, RandomGatheringDataset):
            datasets, weights = datasets.datasets, datasets.weights
        else:
            datasets, weights = parse_item_weight_pairs(datasets)
        datasets, weights = zip(*filter(lambda item: item[1] > 0, zip(datasets, weights)))

        if state is None:
            self._iterators_index = list(range(len(datasets)))
            self._iterators = [iter(dataset) for dataset in datasets]
            self._weights = np.asarray(weights, dtype=np.float64)
            self._weights /= np.sum(self._weights)
            self._random = np.random.default_rng(random_seed)
        else:
            self._iterators_index = list(state['iterators_index'])
            self._iterators = [resume(datasets[index], state)
                               for index, state in zip(self._iterators_index, state["iterators_state"])]
            self._weights = np.asarray(weights)[list(self._iterators_index)]
            self._weights /= np.sum(self._weights)
            self._random = np.random.default_rng(random_seed)
            self._random.bit_generator.state = state['random_state']

    def __next__(self) -> AnySample:
        if len(self._iterators) == 0:
            raise StopIteration

        if len(self._iterators) == 1:
            return next(self._iterators[0])

        while True:
            index = self._random.choice(len(self._iterators), p=self._weights)
            try:
                return next(self._iterators[index])
            except StopIteration:
                self._iterators.pop(index)
                self._iterators_index.pop(index)
                if len(self._iterators) == 1:
                    break
                self._weights = np.delete(self._weights, index)
                self._weights /= np.sum(self._weights)

        return next(self._iterators[0])

    def __pause__(self) -> RandomGatheringIteratorState:
        return RandomGatheringIteratorState(
            iterators_index=tuple(self._iterators_index),
            iterators_state=tuple(pause(iterator) for iterator in self._iterators),
            random_state=self._random.bit_generator.state)


class RandomGatheringDataset(Dataset[AnySample]):
    def __init__(self,
        datasets: Mapping[Iterable[AnySample], float] |
                  Iterable[tuple[Iterable[AnySample], float]] |
                  Iterable[Iterable[AnySample]], *,
        random_seed: int = 42,
    ):
        self._datasets, weights = parse_item_weight_pairs(datasets)
        self._weights = np.asarray(weights)
        self._random_seed = random_seed

    @property
    def datasets(self) -> tuple[Iterable[AnySample], ...]:
        return self._datasets

    @property
    def weights(self) -> np.ndarray:
        return self._weights

    @property
    def random_seed(self) -> int:
        return self._random_seed

    @cached_property
    def __len__(self) -> int:
        # noinspection PyTypeChecker
        return sum(len(dataset) for dataset in self._datasets)

    def __iter__(self) -> RandomGatheringIterator[AnySample]:
        return RandomGatheringIterator(self, None, self._random_seed)

    def __resume__(self,
        state: RandomGatheringIteratorState | None = None, /
    ) -> RandomGatheringIterator[AnySample]:
        return RandomGatheringIterator(self, state, self._random_seed)


class RandomDispatchingWriter(DatasetWriter[AnySample]):
    def __init__(self,
        writers: Mapping[DatasetWriter[AnySample], float] |
                 Iterable[tuple[DatasetWriter[AnySample], float]] |
                 Iterable[DatasetWriter[AnySample]], *,
        random_seed: int = 42,
    ):
        self._writers, weights = parse_item_weight_pairs(writers)
        self._weights = np.asarray(weights) / np.sum(weights)
        self._random = np.random.default_rng(random_seed)

    @property
    def writers(self) -> tuple[DatasetWriter[AnySample], ...]:
        return self._writers

    @property
    def weights(self) -> np.ndarray:
        return self._weights

    def write(self, sample: AnySample):
        writer_i = self._random.choice(len(self._writers), p=self._weights)
        self._writers[writer_i].write(sample)

    def flush(self):
        for writer in self.writers:
            writer.flush()

    def close(self):
        for writer in self.writers:
            writer.close()
